
#ifndef AFTERMATH_NUMERICS_SSSE_MATH_HPP_INCLUDED
#define AFTERMATH_NUMERICS_SSSE_MATH_HPP_INCLUDED

#include "inline.hpp"

#include <mmintrin.h>
#include <smmintrin.h>
#include <xmmintrin.h>

namespace aftermath
{
    namespace numerics
    {
        OP_INLINE __m128 operator +(const __m128& left, const __m128& right) { return _mm_add_ps(left, right); }
        OP_INLINE __m128 operator -(const __m128& left, const __m128& right) { return _mm_sub_ps(left, right); }
        OP_INLINE __m128 operator *(const __m128& left, const __m128& right) { return _mm_mul_ps(left, right); }
        OP_INLINE __m128 operator /(const __m128& left, const __m128& right) { return _mm_div_ps(left, right); }
    
        OP_INLINE __m128& operator +=(__m128& left, const __m128& right) { return left = _mm_add_ps(left, right); }
        OP_INLINE __m128& operator -=(__m128& left, const __m128& right) { return left = _mm_sub_ps(left, right); }
        OP_INLINE __m128& operator *=(__m128& left, const __m128& right) { return left = _mm_mul_ps(left, right); }
        OP_INLINE __m128& operator /=(__m128& left, const __m128& right) { return left = _mm_div_ps(left, right); }


        template <int nRows, int nColumns>
        OP_INLINE void quad_row_transpose(
            const __m128& row0, const __m128& row1, const __m128& row2, const __m128& row3,
            __m128& col0, __m128& col1, __m128& col2, __m128& col3,
            __m128& tmp0, __m128& tmp1, __m128& tmp2, __m128& tmp3, __m128& zero) 
        {
            if (nRows == 0 || nColumns == 0) return;
            if (nRows == 1 && nColumns == 1)
            {
                col0 = row0;
                return;
            }

            // most significant <------ least significant
            // row0: a3 a2 a1 a0
            // row1: b3 b2 b1 b0
            // row2: c3 c2 c1 c0
            // row3: d3 d2 d1 d0

            tmp0 = _mm_shuffle_ps(row0, nRows > 1 ? row1 : zero, 0x44); // tmp0: b1 b0 a1 a0
            if (nColumns > 2) tmp2 = _mm_shuffle_ps(row0, nRows > 1 ? row1 : zero, 0xEE); // tmp2: b3 b2 a3 a2
            if (nRows > 2)
            {
                tmp1 = _mm_shuffle_ps(row2, nRows > 3 ? row3 : zero, 0x44); // tmp1: d1 d0 c1 c0
                if (nColumns > 2) tmp3 = _mm_shuffle_ps(row2, nRows > 3 ? row3 : zero, 0xEE); // tmp3: d3 d2 d3 d2
            }
        
            if (nColumns > 3) col3 = _mm_shuffle_ps(tmp2, nRows > 2 ? tmp3 : zero, 0xDD); // col3: d3 c3 b3 a3
            if (nColumns > 2) col2 = _mm_shuffle_ps(tmp2, nRows > 2 ? tmp3 : zero, 0x88); // col2: d2 c2 b2 a2
            if (nColumns > 1) col1 = _mm_shuffle_ps(tmp0, nRows > 2 ? tmp1 : zero, 0xDD); // col1: d1 c1 b1 a1
            col0 = _mm_shuffle_ps(tmp0, nRows > 2 ? tmp1 : zero, 0x88); // col0: d0 c0 b0 a0
        }

        template <int nRows>
        OP_INLINE void quad_row_transpose_self(
            __m128& row0, __m128& row1, __m128& row2, __m128& row3,
            __m128& tmp0, __m128& tmp2, __m128& zero) 
        {
            if (nRows == 0) return;
            if (nRows == 1) return;

            // most significant <------ least significant
            // row0: a3 a2 a1 a0
            // row1: b3 b2 b1 b0
            // row2: c3 c2 c1 c0
            // row3: d3 d2 d1 d0
        
            tmp0 = _mm_shuffle_ps(row0, nRows > 1 ? row1 : zero, 0x44); // tmp0: b1 b0 a1 a0
            if (nRows > 2)
            {
                tmp2 = _mm_shuffle_ps(row0, row1, 0xEE); // tmp2: b3 b2 a3 a2
                // We don't need row0 and row1 any more; might as well use them
                row0 = _mm_shuffle_ps(row2, nRows > 3 ? row3 : zero, 0x44); // row0: d1 d0 c1 c0
                // row0 (not row1!) now stores what used to be tmp1
                row1 = _mm_shuffle_ps(row2, nRows > 3 ? row3 : zero, 0xEE); // row1: d3 d2 d3 d2
                // row1 now stores what used to be tmp3, we will get rid of it soon
                // Now we don't need original row2 and row3, but we can't overwrite row0 or row1 just yet; so start with the former
            }
        
            if (nRows > 3) row3 = _mm_shuffle_ps(tmp2, row1, 0xDD); // row3: d3 c3 b3 a3
            if (nRows > 2) row2 = _mm_shuffle_ps(tmp2, row1, 0x88); // row2: d2 c2 b2 a2
            // now we don't need row1 (former tmp3): overwrite it
            if (nRows > 1) row1 = _mm_shuffle_ps(tmp0, nRows > 2 ? row0 : zero, 0xDD); // row1: d1 c1 b1 a1
            // finally overwrite row0
            row0 = _mm_shuffle_ps(tmp0, nRows > 2 ? row0 : zero, 0x88); // row0: d0 c0 b0 a0
        }

        template <int nRows>
        OP_INLINE void quad_coupled_row_transpose(
            __m128& urow0, __m128& urow1, __m128& urow2, __m128& urow3,
            __m128& lrow0, __m128& lrow1, __m128& lrow2, __m128& lrow3,
            __m128& utmp0, __m128& utmp1, __m128& utmp2, __m128& utmp3,
            __m128& ltmp0, __m128& ltmp1, __m128& ltmp2, __m128& ltmp3, __m128& zero) 
        {
            if (nRows == 0) return;
            // For upper: nRows = 4, nColumns = nRows
            // For lower: nRows = nRows, nColumns = 4
        
            utmp0 = _mm_shuffle_ps(urow0, urow1, 0x44);
            if (nRows > 2) utmp2 = _mm_shuffle_ps(urow0, urow1, 0xEE);
            utmp1 = _mm_shuffle_ps(urow2, urow3, 0x44);
            if (nRows > 2) utmp3 = _mm_shuffle_ps(urow2, urow3, 0xEE);
        
            quad_row_transpose<nRows, 4>(
                lrow0, lrow1, lrow2, lrow3,
                urow0, urow1, urow2, urow3,
                ltmp0, ltmp1, ltmp2, ltmp3, zero
                );
            //ltmp0 = _mm_shuffle_ps(lrow0, nRows == 1 ? zero : lrow1, 0x44);
            //ltmp2 = _mm_shuffle_ps(lrow0, nRows == 1 ? zero : lrow1, 0xEE);
            //if (nRows > 2)
            //{
            //    ltmp1 = _mm_shuffle_ps(lrow2, nRows == 3 ? zero : lrow3, 0x44);
            //    ltmp3 = _mm_shuffle_ps(lrow2, nRows == 3 ? zero : lrow3, 0xEE);
            //}

            //urow0 = _mm_shuffle_ps(ltmp0, nRows < 3 ? zero : ltmp1, 0x88);
            //urow1 = _mm_shuffle_ps(ltmp0, nRows < 3 ? zero : ltmp1, 0xDD);
            //urow2 = _mm_shuffle_ps(ltmp2, nRows < 3 ? zero : ltmp3, 0x88);
            //urow3 = _mm_shuffle_ps(ltmp2, nRows < 3 ? zero : ltmp3, 0xDD);
        
            lrow0 = _mm_shuffle_ps(utmp0, utmp1, 0x88);
            if (nRows > 1) lrow1 = _mm_shuffle_ps(utmp0, utmp1, 0xDD);
            if (nRows > 2) lrow2 = _mm_shuffle_ps(utmp2, utmp3, 0x88);
            if (nRows > 3) lrow3 = _mm_shuffle_ps(utmp2, utmp3, 0xDD);
        }

        template <int nRows, int nColumns, bool doIncrement>
        OP_INLINE void quad_multiply(
            const __m128& row0, const __m128& row1, const __m128& row2, const __m128& row3,
            const __m128& col0, const __m128& col1, const __m128& col2, const __m128& col3,
            __m128& res0, __m128& res1, __m128& res2, __m128& res3,
            __m128& tmp0, __m128& tmp1, __m128& tmp2, __m128& tmp3, __m128& zero) 
        {
            if (nRows == 0 || nColumns == 0) return;

            // We don't need tmp* then!
            if (nColumns == 1)
            {
                if (doIncrement)
                {
                    res0 += _mm_dp_ps(row0, col0, 0xF1);
                    if (nRows > 1) res1 += _mm_dp_ps(row1, col0, 0xF1);
                    if (nRows > 2) res2 += _mm_dp_ps(row2, col0, 0xF1);
                    if (nRows > 3) res3 += _mm_dp_ps(row3, col0, 0xF1);
                }
                else
                {
                    res0 = _mm_dp_ps(row0, col0, 0xF1);
                    if (nRows > 1) res1 = _mm_dp_ps(row1, col0, 0xF1);
                    if (nRows > 2) res2 = _mm_dp_ps(row2, col0, 0xF1);
                    if (nRows > 3) res3 = _mm_dp_ps(row3, col0, 0xF1);
                }
            }
            else
            {
                tmp0 = _mm_dp_ps(row0, col0, 0xF1);
                tmp1 = _mm_dp_ps(row0, col1, 0xF2);
                if (nColumns > 2) tmp2 = _mm_dp_ps(row0, col2, 0xF4);
                if (nColumns > 3) tmp3 = _mm_dp_ps(row0, col3, 0xF8);

                if (doIncrement)
                {
                    res0 += tmp0;
                    res0 += tmp1;
                }
                else res0 = tmp0 + tmp1;
                if (nColumns > 2) res0 += tmp2;
                if (nColumns > 3) res0 += tmp3;

                if (nRows > 1)
                {
                    tmp0 = _mm_dp_ps(row1, col0, 0xF1);
                    tmp1 = _mm_dp_ps(row1, col1, 0xF2);
                    if (nColumns > 2) tmp2 = _mm_dp_ps(row1, col2, 0xF4);
                    if (nColumns > 3) tmp3 = _mm_dp_ps(row1, col3, 0xF8);
            
                    if (doIncrement)
                    {
                        res1 += tmp0;
                        res1 += tmp1;
                    }
                    else res1 = tmp0 + tmp1;
                    if (nColumns > 2) res1 += tmp2;
                    if (nColumns > 3) res1 += tmp3;
                }
                if (nRows > 2)
                {
                    tmp0 = _mm_dp_ps(row2, col0, 0xF1);
                    tmp1 = _mm_dp_ps(row2, col1, 0xF2);
                    if (nColumns > 2) tmp2 = _mm_dp_ps(row2, col2, 0xF4);
                    if (nColumns > 3) tmp3 = _mm_dp_ps(row2, col3, 0xF8);
            
                    if (doIncrement)
                    {
                        res2 += tmp0;
                        res2 += tmp1;
                    }
                    else res2 = tmp0 + tmp1;
                    if (nColumns > 2) res2 += tmp2;
                    if (nColumns > 3) res2 += tmp3;
                }
                if (nRows > 3)
                {
                    tmp0 = _mm_dp_ps(row3, col0, 0xF1);
                    tmp1 = _mm_dp_ps(row3, col1, 0xF2);
                    if (nColumns > 2) tmp2 = _mm_dp_ps(row3, col2, 0xF4);
                    if (nColumns > 3) tmp3 = _mm_dp_ps(row3, col3, 0xF8);
            
                    if (doIncrement)
                    {
                        res3 += tmp0;
                        res3 += tmp1;
                    }
                    else res3 = tmp0 + tmp1;
                    if (nColumns > 2) res3 += tmp2;
                    if (nColumns > 3) res3 += tmp3;
                }
            }
        }
    }
}

#endif // AFTERMATH_NUMERICS_SSSE_MATH_HPP_INCLUDED
